/*
 * OdeSolverImp.cpp
 *
 *  Created on: 28 Jul 2011
 *      Author: Allan
 */

#include "Numerics/OdeSolver.h"

// C++ includes
#include <iostream>
#include <functional>
using namespace std;
using namespace std::placeholders;

// GSL includes
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_odeiv2.h>

class OdeSolverImp
{
public:
	OdeSolverImp();
	
	OdeSolverImp(unsigned dimension, OdeAlgorithm algorithm);
	
	~OdeSolverImp();
	
	void Initialiaze(unsigned dimension, OdeAlgorithm algorithm);
	
	void SetAlgorithm(OdeAlgorithm algorithm);
	
	void SetControlStandard(double epsAbs, double epsRel, double yFactor, double dydtFactor);
	
	void SetControlSolution(double epsAbs, double epsRel);
	
	void SetControlSolutionDerivative(double epsAbs, double epsRel);
	
	bool Iterate(OdeProblem& problem, double& t, double tEnd, double& dt, VectorXd& y);
	
	bool Solve(OdeProblem& problem, double& t, double tEnd, double& dt, VectorXd& y);

public:
	struct Params
	{
		OdeProblem* problem;
		VectorXd y;
		VectorXd f;
		VectorXd dfdt;
		MatrixXd dfdy;
	};
	
private:
	unsigned dimension;
	
	Params params;
	
	const gsl_odeiv2_step_type* type;
	
	gsl_odeiv2_step* step; 
	
	gsl_odeiv2_control* control;
	
	gsl_odeiv2_evolve* evolve;
	
	double* yGSL;
};

OdeSolver::OdeSolver()
{
	imp = std::unique_ptr<OdeSolverImp>(new OdeSolverImp());
}

OdeSolver::OdeSolver(unsigned dimension, OdeAlgorithm algorithm)
{
	imp = std::unique_ptr<OdeSolverImp>(new OdeSolverImp(dimension, algorithm));
}

OdeSolver::~OdeSolver()
{}

void OdeSolver::Initialiaze(unsigned dimension, OdeAlgorithm algorithm)
{
	imp->Initialiaze(dimension, algorithm);
}

void OdeSolver::SetControlStandard(double epsAbs, double epsRel, double yFactor, double dydtFactor)
{
	imp->SetControlStandard(epsAbs, epsRel, yFactor, dydtFactor);
}

void OdeSolver::SetControlSolution(double epsAbs, double epsRel)
{
	imp->SetControlSolution(epsAbs, epsRel);
}

void OdeSolver::SetControlSolutionDerivative(double epsAbs, double epsRel)
{
	imp->SetControlSolutionDerivative(epsAbs, epsRel);
}

bool OdeSolver::Iterate(OdeProblem& problem, double& t, double tEnd, double& dt, VectorXd& y)
{
	return imp->Iterate(problem, t, tEnd, dt, y);
}

bool OdeSolver::Solve(OdeProblem& problem, double& t, double tEnd, double& dt, VectorXd& y)
{
	return imp->Solve(problem, t, tEnd, dt, y);
}

int GSLFunction(double t, const double y[], double f[], void* params);
int GSLJacobian(double t, const double y[], double* dfdy, double dfdt[], void* params);

OdeSolverImp::OdeSolverImp() :
dimension(0), type(gsl_odeiv2_step_rk4), step(0), control(0), evolve(0), yGSL(0)
{}

OdeSolverImp::OdeSolverImp(unsigned dimension, OdeAlgorithm algorithm) :
dimension(0), type(gsl_odeiv2_step_rk4), step(0), control(0), evolve(0), yGSL(0)
{
	Initialiaze(dimension, algorithm);
}

OdeSolverImp::~OdeSolverImp()
{
	gsl_odeiv2_evolve_free(evolve);
	gsl_odeiv2_control_free (control);
	gsl_odeiv2_step_free(step);
	
	delete [] yGSL;
}

void OdeSolverImp::Initialiaze(unsigned dimension, OdeAlgorithm algorithm)
{
	// Set the data member (dimension)
	this->dimension = dimension;
	
	// Detemine the (type) of the GSL ode algorithm 
	switch(algorithm)
	{
	case RungeKutta2:
		type = gsl_odeiv2_step_rk2;
		break;
	case RungeKutta4:
		type = gsl_odeiv2_step_rk4;
		break;
	case RungeKuttaFehlberg45:
		type = gsl_odeiv2_step_rkf45;
		break;
	case RungeKuttaCashKarp45:
		type = gsl_odeiv2_step_rkck;
		break;
	case RungeKuttaPrinceDormand8:
		type = gsl_odeiv2_step_rk8pd;
		break;
	case RungeKuttaImplicitGaussian1:
		type = gsl_odeiv2_step_rk1imp;
		break;
	case RungeKuttaImplicitGaussian2:
		type = gsl_odeiv2_step_rk2imp;
		break;
	case RungeKuttaImplicitGaussian4:
		type = gsl_odeiv2_step_rk4imp;
		break;
	case BulirschStoerImplicit:
		type = gsl_odeiv2_step_bsimp;
		break;
	case MultistepAdams:
		type = gsl_odeiv2_step_msadams;
		break;
	case MultistepBDF:
		type = gsl_odeiv2_step_msbdf;
		break;
	default:
		type = gsl_odeiv2_step_rk4;
		break;
	}
	
	// Set the data member (params)
	params.problem = 0; // This field is set in the methods Iterate and Solve
	params.y       = VectorXd::Zero(dimension);
	params.f       = VectorXd::Zero(dimension);
	params.dfdt    = VectorXd::Zero(dimension);
	params.dfdy    = MatrixXd::Zero(dimension, dimension);
	
	// Free the GSL objects (step), (evolve) and (yGSL) if they have already been initialized
	if(step   != 0) gsl_odeiv2_step_free(step);
	if(evolve != 0) gsl_odeiv2_evolve_free(evolve);
	if(yGSL   != 0) delete [] yGSL;
	
	// Initialiaze the GSL objects (step), (evolve) and (yGSL) 
	step   = gsl_odeiv2_step_alloc(type, dimension);
	evolve = gsl_odeiv2_evolve_alloc(dimension);
	yGSL   = new double[dimension];
}

void OdeSolverImp::SetControlStandard(double epsAbs, double epsRel, double ay, double adydt)
{
	if(control != 0) gsl_odeiv2_control_free(control);
	
	control = gsl_odeiv2_control_standard_new(epsAbs, epsRel, ay, adydt);
}
	
void OdeSolverImp::SetControlSolution(double epsAbs, double epsRel)
{
	if(control != 0) gsl_odeiv2_control_free(control);
	
	control = gsl_odeiv2_control_standard_new(epsAbs, epsRel, 1.0, 0.0);
}

void OdeSolverImp::SetControlSolutionDerivative(double epsAbs, double epsRel)
{
	if(control != 0) gsl_odeiv2_control_free(control);
	
	control = gsl_odeiv2_control_standard_new(epsAbs, epsRel, 0.0, 1.0);
}

bool OdeSolverImp::Iterate(OdeProblem& problem, double& t, double tEnd, double& dt, VectorXd& y)
{
	// Set the (problem) field of the instance (params) 
	params.problem = &problem;
	
	// The GSL variable that stores the function, jacobian, etc.
	gsl_odeiv2_system sys = {GSLFunction, GSLJacobian, dimension, &params};
	
	// Transfer the data in the given argument (y) to (yGSL)
	for(unsigned i = 0; i < dimension; ++i) yGSL[i] = y[i];
	
	// Apply the iteration method
	int status = gsl_odeiv2_evolve_apply(evolve, control, step, &sys, &t, tEnd, &dt, yGSL);
	
	// Check if the iteration failed
	if(status != GSL_SUCCESS) return false;
		
	// Transfer the data in (yGSL) to the given argument (y) 
	for(unsigned i = 0; i < dimension; ++i) y[i] = yGSL[i];
	
	return true;
}

bool OdeSolverImp::Solve(OdeProblem& problem, double& t, double tEnd, double& dt, VectorXd& y)
{
	// Set the (problem) field of the instance (params) 
	params.problem = &problem;
	
	// The GSL variable that stores the function, jacobian, etc.
	gsl_odeiv2_system sys = {GSLFunction, GSLJacobian, dimension, &params};
	
	// Transfer the data in the given argument (y) to (yGSL)
	for(unsigned i = 0; i < dimension; ++i) yGSL[i] = y[i];
	
	// Loop until (t) reaches (tEnd)
	while(t < tEnd)
	{
		// Apply the iteration method
		int status = gsl_odeiv2_evolve_apply(evolve, control, step, &sys, &t, tEnd, &dt, yGSL);
		
		// Check if the iteration failed
		if(status != GSL_SUCCESS) return false;
	}
	
	// Transfer the data in (yGSL) to the given argument (y) 
	for(unsigned i = 0; i < dimension; ++i) y[i] = yGSL[i];
	
	return true;
}

int GSLFunction(double t, const double y[], double f[], void* params)
{
	// Cast (params) to the correct type
	OdeSolverImp::Params& p = *((OdeSolverImp::Params*)params);
	
	// The dimension of the array (y)
	const unsigned dimension = p.y.rows();
	
	// Transfer the data in (y) to (p.y)
	for(unsigned i = 0; i < dimension; ++i) p.y[i] = y[i];
	
	// Compute the Eigen vector (p.f)
	p.problem->Function(t, p.y, p.f);
	
	// Transfer the data in (p.f) to (f)
	for(unsigned i = 0; i < dimension; ++i) f[i] = p.f[i];
	
	return GSL_SUCCESS;
}

int GSLJacobian(double t, const double y[], double* dfdy, double dfdt[], void* params)
{
	// Cast (params) to the correct type
	OdeSolverImp::Params& p = *((OdeSolverImp::Params*)params);
	
	// The dimension of the array (y)
	const unsigned dimension = p.y.rows();
	
	// Transfer the data in (p.y) to (y)
	for(unsigned i = 0; i < dimension; ++i) p.y[i] = y[i];
	
	// Compute the Eigen matrix (p.dfdy) and Eigen vector (p.dfdt) 
	p.problem->Jacobian(t, p.y, p.dfdy, p.dfdt);
	
	// Transfer the data in (p.dfdy) to (dfdy)
	for(unsigned i = 0; i < dimension; ++i)
		for(unsigned j = 0; j < dimension; ++j)
			dfdy[i + dimension * j] = p.dfdy(i, j);
	
	// Transfer the data in (p.dfdt) to (dfdt)
	for(unsigned i = 0; i < dimension; ++i)
		dfdt[i] = p.dfdt[i];
	
	return GSL_SUCCESS;
}
